# Copyright (c) Microsoft Corporation.
# Licensed under the MIT License.

import json

from data_formulator.agents.agent_utils import extract_json_objects, generate_data_summary, extract_code_from_gpt_response
from data_formulator.agents.agent_data_transform_v2 import completion_response_wrapper

import data_formulator.py_sandbox as py_sandbox

import traceback


import logging

logger = logging.getLogger(__name__)


SYSTEM_PROMPT = '''您是一位帮助用户推荐可视化数据的数据科学家。
用户将提供他们想要创建的可视化信息, 您的任务是推荐可用于创建可视化的转换数据, 并编写 Python 函数来实现数据转换。
推荐和转换函数应基于用户提供的 [CONTEXT] 和 [GOAL]。
[CONTEXT] 显示当前数据集的情况, [GOAL] 描述用户的数据用途。

**重要规则:**
- 严禁对个人的性别、生物性别、性取向、宗教、种族、国籍、民族、政治立场、社会经济地位、心理健康、隐形残疾、医疗状况、人格类型、社会印象、情绪状态和认知状态做出任何假设或判断。
- 严禁创建可能用于年龄歧视的计算公式。任何形式的年龄歧视 (显性或隐性) 都被严格禁止。
- 如遇上述问题, 请用 np.nan 填充对应列。

具体而言, 您需要基于 [CONTEXT] 和 [GOAL] 分两步完成, 第 1 步即推断适用的数据结构, 第 2 步则是在 [OUTPUT] 部分创建对应的 Python 函数。

    1. 根据用户 [GOAL] 创建表示推断意图的 JSON 对象, 格式如下:

{
    "mode": "" // string 类型, 可选项有 "infer", "overview", "distribution", "summary"
    "recommendation": "..." // string 类型, 说明推荐理由
    "output_fields": [...] // string[] 类型, 描述输出数据应包含的目标字段 (i.e., 转换后数据的最终目标), 建议在此保留中间字段
    "chart_type": "" // string 类型, 可选项有 "point", "bar", "line", "boxplot". "chart_type" 需从用户指令推断或自行推荐
    "visualization_fields": [] // string[] 类型: 从 "output_fields" 中选择需可视化的子集 (不超过 3 个), 按图表 "x, y, 图例" 顺序排列，排除中间字段
}

具体步骤:
    (1) 如果用户的 [GOAL] 明确, 直接推断用户意图。将 "mode" 设为 "infer", 并根据用户描述创建 "output_fields" 和 "visualization_fields"。
    (2) 如果用户的 [GOAL] 不明确, 向用户推荐以下方案:
        - "mode" 的可选项有 "distribution", "overview", "summary", 需选择其一:
            * 若选择 "overview", 并且数据是宽格式 (wide format), 则将其重塑为长格式 (long format)。
            * 若选择 "distribution", 则选择若干有可视化潜力的字段进行联合展示。
            * 若选择 "summary", 则计算聚合、统计以展示数据特征。
        - 在 "recommendation" 中描述推荐理由。
        - 根据推荐确定理想的输出数据, 输出数据必须是整洁数据格式 (tidy format)。
        - 推荐应被可视化的字段组合。
    (3) "visualization_fields" 需按字段在 "x 轴 / y 轴 / 图例" 中的使用顺序排列, 不包含 "output_fields" 中的中间字段。
    (4) "visualization_fields" 最多包含 3 个字段 (对应 x / y / 图例)。
    (5) "chart_type" 必须是指定图表类型之一 ("point" / "bar" / "line" / "boxplot")。

    2. 根据推断目标编写 Python 函数, 函数输入为 dataframe ("df") (或基于 [CONTEXT] 中表格的多个 dataframs), 输出为转换后的 dataframe ("transformed_df"), "transformed_df" 应包含优化目标中所有 "output_fields"。
Python 函数必须遵循 [TEMPLATE] 提供的模板, 不要导入其他库或修改函数名称。函数应尽可能简洁易读。
若基于 "output_fields" 无需数据转换, 转换函数可直接返回原数据框 ("return df")。

[TEMPLATE]

```python
import pandas as pd
import collections
import numpy as np

def transform_data(df1, df2, ...):
    # 在此完成模板代码
    return transformed_df
```

注意事项:
- 若用户提供单个表格, 函数应定义为 def transform_data(df1); 若提供多个表格, 应定义为 def transform_data(df1, df2, ...) 并需考虑表连接逻辑。
- 建议使用表名指代输入的 dataframes。例如: 若用户提供 city 和 weather 两个表, 可使用 transform_data(df_city, df_weather) 进行引用。

    3. [OUTPUT] 必须仅包含表示优化目标的 JSON 对象和转换代码的 Python 代码块, 禁止添加额外文本解释。
'''

example = """
For example:

[CONTEXT]

数据集、字段摘要及样例数据:

table_0 (student_exam) fields:
	student -- type:  int64, values:  1,  2,  3, ..., 997, 998, 999, 1000
	  major -- type: object, values: liberal arts, science
	   math -- type:  int64, values:  0,  8, 18, ...,  97,  98,  99,  100
	reading -- type:  int64, values: 17, 23, 24, ...,  96,  97,  99,  100
	writing -- type:  int64, values: 10, 15, 19, ...,  97,  98,  99,  100

table_0 (student_exam) sample:

```
|student|major|math|reading|writing
0|1|liberal arts|72|72|74
1|2|liberal arts|69|90|88
2|3|liberal arts|90|95|93
3|4|science|47|57|44
4|5|science|76|78|75
......
```

[GOAL]

{"goal": "根据 "平均分" 对学生进行排名"}

[OUTPUT]

{
    "mode": "infer",
    "recommendation": "要根据平均分排名, 则需要计算每个学生的 "平均分", 然后根据 "平均分" 排序数据并分配排名。",
    "output_fields": ["student", "major", "平均分", "排名"],
    "visualization_fields": ["student", "平均分"],
}

```python
import pandas as pd
import collections
import numpy as np

def transform_data(df):
    df['平均分'] = df[['math', 'reading', 'writing']].mean(axis=1)
    df = df.sort_values(by='平均分', ascending=False)
    df['排名'] = df['平均分'].rank(ascending=False, method='dense').astype(int)
    transformed_df = df[['student', 'major', '平均分', '排名']]
    return transformed_df
```
"""

class DataRecAgent(object):

    def __init__(self, client, system_prompt=None):
        self.client = client
        self.system_prompt = system_prompt if system_prompt is not None else SYSTEM_PROMPT

    def process_gpt_response(self, input_tables, messages, response):
        """处理 GPT 响应并执行代码"""

        #log = {'messages': messages, 'response': response.model_dump(mode='json')}

        if isinstance(response, Exception):
            result = {'status': 'other error', 'content': str(response.body)}
            return [result]

        candidates = []
        for choice in response.choices:

            logger.info("\n=== Data recommendation result ===>\n")
            logger.info(choice.message.content + "\n")

            json_blocks = extract_json_objects(choice.message.content + "\n")
            if len(json_blocks) > 0:
                refined_goal = json_blocks[0]
            else:
                refined_goal = { 'mode': "", 'recommendation': "", 'output_fields': [], 'visualization_fields': [], }

            code_blocks = extract_code_from_gpt_response(choice.message.content + "\n", "python")

            if len(code_blocks) > 0:
                code_str = code_blocks[-1]

                try:
                    result = py_sandbox.run_transform_in_sandbox2020(code_str, [t['rows'] for t in input_tables])
                    result['code'] = code_str

                    if result['status'] == 'ok':
                        result['content'] = json.loads(result['content'])
                    else:
                        logger.info(result['content'])
                except Exception as e:
                    logger.warning('other error:')
                    error_message = traceback.format_exc()
                    logger.warning(error_message)
                    result = {'status': 'other error', 'code': code_str, 'content': f"Unexpected error: {error_message}"}
            else:
                result = {'status': 'error', 'code': "", 'content': "No code block found in the response. The model is unable to generate code to complete the task."}

            result['dialog'] = [*messages, {"role": choice.message.role, "content": choice.message.content}]
            result['agent'] = 'DataRecAgent'
            result['refined_goal'] = refined_goal
            candidates.append(result)

        logger.info("=== Recommendation Candidates ===>")
        for candidate in candidates:
            for key, value in candidate.items():
                if key in ['dialog', 'content']:
                    logger.info(f"##{key}:\n{str(value)[:1000]}...")
                else:
                    logger.info(f"## {key}:\n{value}")

        return candidates


    def run(self, input_tables, description, n=1):

        data_summary = generate_data_summary(input_tables, include_data_samples=True)

        user_query = f"[CONTEXT]\n\n{data_summary}\n\n[GOAL]\n\n{description}\n\n[OUTPUT]\n"

        logger.info(user_query)

        messages = [{"role":"system", "content": self.system_prompt},
                    {"role":"user","content": user_query}]

        response = completion_response_wrapper(self.client, messages, n)

        return self.process_gpt_response(input_tables, messages, response)


    def followup(self, input_tables, dialog, new_instruction: str, n=1):
        """扩展输入数据 (JSON 格式) 以包含新字段"""

        logger.info(f"GOAL: \n\n{new_instruction}")

        messages = [*dialog, {"role":"user", "content": f"Update: \n\n{new_instruction}"}]

        response = completion_response_wrapper(self.client, messages, n)

        return self.process_gpt_response(input_tables, messages, response)
